import json
from os.path import join
import paddle.fluid.layers as L
import paddle.fluid as F
from paddle import to_tensor
import paddle.fluid.dygraph as D
from ernie.modeling_ernie import ErnieModel
from myReader.nil_types import *
import numpy as np


class ModelWithErnie(D.Layer):
    # 减小计算量
    def __init__(self, args, version_cfg):
        super(ModelWithErnie, self).__init__()
        self.args = args
        self.version_cfg = version_cfg
        self.ernie = ErnieModel.from_pretrained(args.model_name_or_path, num_out_pooler=2)

        self.entity_embed_layer = D.Embedding([len(TYPE2ID), 768])
        # self.word_embed_layer = self.ernie_mt.word_emb

        self.sim_ffn = D.Linear(768, 1)  # sim
        self.nil_ffn = D.Linear(768, len(TYPE2ID))
        self.sim_pos, self.sim_neg = to_tensor([0.0]), to_tensor([0.0])
        self.register_buffer('sim_pos', self.sim_pos)
        self.register_buffer('sim_neg', self.sim_neg)

    def forward(self, inputs: dict):
        if 'label_sim' in inputs:
            label_sim = inputs.pop('label_sim')
            if self.training:
                self.sim_pos += L.cast(label_sim == 0.5, 'float32').sum()
                self.sim_neg += L.cast(label_sim == -0.5, 'float32').sum()
        pooled, _, = self.ernie(**inputs)
        out0 = L.tanh(L.squeeze(self.sim_ffn(pooled[0]), [1]))
        out1 = L.softmax(self.nil_ffn(pooled[1]))
        return out0, out1

    def sim_threshold(self):
        """
        NIL阈值
        x - neg
        """
        if self.sim_pos == self.sim_neg:
            return 0
        prob_pos = self.sim_pos / (self.sim_pos + self.sim_neg + 1e-7)
        prob_neg = 1.0 - prob_pos
        return (prob_pos - prob_neg) * 0.99999
